
#include "CudaManager.h"
#include "Runtime.Low/Services.h"
#include <windows.h>

using namespace ttg::devices;
using namespace ttg_internal::devices;

//----------------------
//--- Cuda-functions ---
//----------------------

static int (__stdcall *ttgCuInit)(unsigned int flags) = NULL;
static int (__stdcall *ttgCuDeviceGet)(int *device, int ordinal) = NULL;
static int (__stdcall *ttgCuDeviceGetCount)(int *count) = NULL;
static int (__stdcall *ttgCuDeviceGetName)(char *name, int len, int dev) = NULL;
static int (__stdcall *ttgCuDeviceComputeCapability)(int *major, int *minor, int dev) = NULL;
static int (__stdcall *ttgCuDeviceTotalMem)(unsigned int *bytes, int dev) = NULL;
static int (__stdcall *ttgCuDeviceGetProperties)(CudaProps *prop, int dev) = NULL;

//-------------------
//--- CudaManager ---
//-------------------

std::vector<std::string> CudaManager::getCudaDlls()
{
	std::vector<std::string> res;
	char buf[4 * 1024];
	//Getting reference to Cuda Toolkit.
	if (GetEnvironmentVariableA("CUDA_BIN_PATH", buf, 4 * 1024))
		res.push_back(std::string(buf) + "\\cudart.dll");
	//Getting reference to "nvcuda.dll".
	if (GetEnvironmentVariableA("windir", buf, 4 * 1024))
		res.push_back(std::string(buf) + "\\System32\\nvcuda.dll");
	return res;
}

void CudaManager::releaseCudaDevices()
{
	for (size_t i = 0; i < cuda_devices.size(); i++)
	{
		cuda_devices[i].second->terminate();
		SAFE_DELETE(cuda_devices[i].second);
		SAFE_DELETE(cuda_devices[i].first);
	}
	cuda_devices.clear();
}

void CudaManager::init(bool)
{
	//Checking.
	if (!ttg::services::DllLoader::getRef()->isStarted())
		return;
	//Getting cuda params.
	locker.enter();
	std::vector<std::string> cuda_dlls = getCudaDlls();
	for (size_t i = 0; i < cuda_dlls.size(); i++)
	{
		try
		{
			//Binding.
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuInit", (FARPROC &)ttgCuInit);
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuDeviceGet", (FARPROC &)ttgCuDeviceGet);
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuDeviceGetCount", (FARPROC &)ttgCuDeviceGetCount);
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuDeviceGetName", (FARPROC &)ttgCuDeviceGetName);
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuDeviceComputeCapability", (FARPROC &)ttgCuDeviceComputeCapability);
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuDeviceTotalMem", (FARPROC &)ttgCuDeviceTotalMem);
			ttg::services::DllLoader::getRef()->load(cuda_dlls[i].c_str(),
													 "cuDeviceGetProperties", (FARPROC &)ttgCuDeviceGetProperties);
			//Loading.
			int dev_count;
			int err_code;
			if ((err_code = ttgCuInit(0)) != 0 ||
				(err_code = ttgCuDeviceGetCount(&dev_count)) != 0)
				break;
			for (int i = 0; i < dev_count; i++)
			{
				int cuda_device;
				char name[1024];
				CudaProps *props = new CudaProps();
				if (ttgCuDeviceGet(&cuda_device, i) != 0 ||
					ttgCuDeviceGetName(name, 1024, cuda_device) != 0 ||
					ttgCuDeviceComputeCapability(&props->major, &props->minor, cuda_device) != 0 ||
					ttgCuDeviceTotalMem(&props->totalGlobalMem, cuda_device) != 0 ||
					ttgCuDeviceGetProperties(props, cuda_device) != 0)
				{
					SAFE_DELETE(props);
					break;
				}
				props->name = name;
				cuda_devices.push_back(std::make_pair(props, new CudaThread()));
			}
			//All is done - returning.
			locker.leave();
			return;
		}
		catch (ttg::Error &) { }
	}
	//Some error has been detected - releasing cuda_devices.
	releaseCudaDevices();
	locker.leave();
}

CudaManager::CudaManager()
{
	if (ttg::services::DllLoader::getRef()->isStarted())
		init(true);
	else
		ttg::services::DllLoader::getRef()->onStart()
			+= ttg::structures::make_delegate(this, &CudaManager::init);
}

void CudaManager::perform(ttg::structures::InvokableObject *action)
{
	for (size_t i = 0; i < cuda_devices.size(); i++)
		if (cuda_devices[i].second->tryAcquire())
		{
			cuda_devices[i].second->perform(action);
			return;
		}
	throw ttg::Error(ttg::Error::FATAL, "CudaManager", "perform",
					 "There aren't any free CUDA devices - wrong task scheduler!");
}

CudaManager::~CudaManager()
{
	locker.enter();
	releaseCudaDevices();
	locker.leave();
}
